package com.upup.train_management.security;

import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.authority.AuthorityUtils;
import org.springframework.security.web.authentication.AuthenticationSuccessHandler;
import org.springframework.stereotype.Component;
import org.springframework.security.web.authentication.SavedRequestAwareAuthenticationSuccessHandler;

import java.io.IOException;
import java.util.Set;

@Component
public class CustomAuthenticationSuccessHandler extends SavedRequestAwareAuthenticationSuccessHandler {
    
    @Override
    public void onAuthenticationSuccess(HttpServletRequest request, 
                                      HttpServletResponse response, 
                                      Authentication authentication) throws IOException {
        
        Set<String> roles = AuthorityUtils.authorityListToSet(authentication.getAuthorities());
        
        // 检查是否有重定向参数
        String redirectParam = request.getParameter("redirect");
        
        if (redirectParam != null && !redirectParam.isEmpty()) {
            // 如果有重定向参数，验证安全性后进行重定向
            if (isValidRedirectUrl(redirectParam)) {
                response.sendRedirect(redirectParam);
                return;
            }
        }
        
        // 根据用户角色重定向到不同页面
        if (roles.contains("ROLE_ADMIN")) {
            response.sendRedirect("/admin/dashboard");
        } else {
            response.sendRedirect("/index");
        }
    }
    
    private boolean isValidRedirectUrl(String url) {
        // 验证重定向URL的安全性
        return url.startsWith("/") && !url.startsWith("//");
    }
}
